from typing import List
from NS import NS
from copy import deepcopy

# Rcon for extend key
RCON =  [0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80, 0x1b, 0x36] 

# Add round key
class AK:
    def __init__(self, key : List[List[int]]) -> None:
        self.keys = [key]
        for i in range(1,11):
            W = self.getTranspose(self.keys[i-1])
            w0 = self.Xor(W[0], self.T(W[3],i))
            w1 = self.Xor(W[1], w0)
            w2 = self.Xor(W[2], w1)
            w3 = self.Xor(W[3], w2)
            self.keys.append(self.getTranspose([w0,w1,w2,w3]))
            

    # get transpose matrix
    def getTranspose(self, nums: List[List[int]]) -> List[List[int]]:
        n = deepcopy(nums)
        return [ [n[y][x] for y in range(len(n))] for x in range(len(n))]

    # print round key (for debug)           
    def printRoundKey(self, round: int) -> None:
        nums = self.getTranspose(self.keys[round])
        count = 4 * round
        for n in nums:
            k = ''
            for s in n:
                s = hex(s)
                # format e.g:  0xa -> 0x0a
                if len(s) == 3:
                    s = s[0:2] + '0' + s[2]
                k += s[2:]
            tag = 'W'+str(count)
            # add blank for alignment
            if count < 10:  tag += ' '
            print(tag,k)
            count += 1

    # get round key
    def getRoundKey(self, round: int) -> List[List[int]]:
        return deepcopy(self.keys[round])

    # RotWord the W[i]
    def rotWord(self, n: List[int]) -> List[int]:
        nums = deepcopy(n)
        t = nums[0]
        nums[0] = nums[1]
        nums[1] = nums[2]
        nums[2] = nums[3]
        nums[3] = t
        return nums

    # subByres the W[i]
    def subBytes(self, n: List[int]) -> List[int]:
        nums = deepcopy(n)
        nums = NS().nsNums(nums)
        return nums
        
    # RCON in round i
    def rconXor(self, n: List[int], i: int) -> List[int]:
        nums = deepcopy(n)
        nums[0] ^= (RCON[i-1])
        return nums

    # n1 xor n2
    def Xor(self, n1: List[int], n2: List[int]) -> List[int]:
        return [n1[i]^n2[i] for i in range(len(n1))]
    
    # function T
    def T(self, n: List[int], i: int) -> List[int]:
        nums = deepcopy(n)
        return self.rconXor(self.subBytes(self.rotWord(nums)), i)
